import pandas as pd
import os
import glob

import datetime as dt
import torch
from torch.utils.data import DataLoader
# from  visdom import Visdom
from tqdm import tqdm
import numpy as np
from sklearn.metrics import r2_score
from scipy.stats import pearsonr
import matplotlib.pyplot as plt

from model import get_crnn,get_res
from dataset import SunDataset
from utils import AverageMeter

year = 2019
days_delay = 3 # 3日延迟或者是4日延迟
data_root = 'D:/datasets/SunData'

def main():
    if not os.path.exists('logs/'):
        os.makedirs('logs/')
    # vis = # visdom(env='sun_pred', log_to_filename='logs/sun_pred_{}.log'.format(
    #    dt.datetime.now().strftime('%Y-%m-%d_T%H-%M-%S')))
    # assert  vis.check_connection()

    device = 'cuda:0'
    checkpoint = torch.load('saved_models/{}days.pth'.format(days_delay))['model']
    # checkpoint = torch.load('saved_models/res18_sig_best.pth')['model']
    model = get_res()
    model.load_state_dict(checkpoint)
    model.to(device)

    criterion = torch.nn.L1Loss()
    # criterion = torch.nn.MSELoss()
    loss_meter = AverageMeter()

    pred_set = SunDataset(root=data_root, split='test', year=year, size=224, gray=False, delay=days_delay)
    pred_loader = DataLoader(pred_set, batch_size=1,
                             shuffle=False, num_workers=0)

    gts = []
    preds = []

    model.eval()
    bar = tqdm(total=len(pred_loader))
    with torch.no_grad():
        for i, (x, y) in enumerate(pred_loader):

            y_hat = model(x.to(device)).cpu()
            y_hat = y_hat*700+200
            gts.append(y.item())
            preds.append(y_hat.flatten().item())

            loss = criterion(y_hat.flatten(), y)
            loss_meter.update(loss.item())

            # vis.line([[y, y_hat]], [i], win='pred', opts={
            #         'legend': ['gt', 'pred']}, update='append')
            # vis.line([loss_meter.avg], [i], win='loss', update='append')

            bar.update()
            bar.set_description('loss:{:.5f}'.format(loss_meter.avg))

        bar.close()
    plt.figure()
    plt.plot(gts)
    plt.plot(preds)
    plt.legend(['gts', 'preds'])
    plt.figure()
    plt.scatter(gts,preds)
    plt.xlabel('gts')
    plt.ylabel('preds')
    gts = np.array(gts)
    preds = np.array(preds)

    p_sorce = pearsonr(gts, preds)
    print("Pearsonr score of :{:0.4f} and {:.4f}".format(p_sorce[0], p_sorce[1]))
    # print("R1 score of :{:0.4f}".format(np.corrcoef(gts,preds)[0,1]))
    print("The R2 score on the Test set is:\t{:0.3f}".format(
        r2_score(gts, preds)))
    plt.show()

if __name__ == "__main__":
    main()
